{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# RAG with Feast Feature Store"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Install required dependencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pip install --quiet feast[milvus] sentence-transformers datasets\n",
    "%pip install bigtree==0.19.2\n",
    "%pip install marshmallow==3.10.0 "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Retrieve test dataset and chunk it"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from datasets import load_dataset\n",
    "# load wikipedia dataset - 1% of the training split\n",
    "dataset = load_dataset(\n",
    "    \"facebook/wiki_dpr\",\n",
    "    \"psgs_w100.nq.exact\",\n",
    "    split=\"train[:1%]\",\n",
    "    with_index=False,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def chunk_dataset(examples, chunk_size=100, overlap=20, max_chars=500):\n",
    "    all_chunks = []\n",
    "    all_ids = []\n",
    "    all_titles = []\n",
    "\n",
    "    for i, text in enumerate(examples['text']):  # Iterate over texts in the batch\n",
    "        words = text.split()\n",
    "        chunks = []\n",
    "        for j in range(0, len(words), chunk_size - overlap):\n",
    "            chunk_words = words[j:j + chunk_size]\n",
    "            if len(chunk_words) < 20:\n",
    "                continue\n",
    "            chunk_text_value = ' '.join(chunk_words)  # Store the chunk text\n",
    "            chunk_text_value = chunk_text_value[:max_chars]\n",
    "            chunks.append(chunk_text_value)\n",
    "            all_ids.append(f\"{examples['id'][i]}_{j}\")  # Unique ID for the chunk\n",
    "            all_titles.append(examples['title'][i])\n",
    "\n",
    "        all_chunks.extend(chunks)\n",
    "\n",
    "    return {'id': all_ids, 'title': all_titles, 'text': all_chunks}\n",
    "\n",
    "\n",
    "chunked_dataset = dataset.map(\n",
    "    chunk_dataset,\n",
    "    batched=True,\n",
    "    remove_columns=dataset.column_names,\n",
    "    num_proc=1\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define embedding model and generate embeddings"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sentence_transformers import SentenceTransformer\n",
    "\n",
    "sentences = chunked_dataset[\"text\"]\n",
    "# Take the first 100 sentences\n",
    "test_sentences = sentences[:100]\n",
    "# load pretrained sentence transformer model and create embeddings\n",
    "embedding_model = SentenceTransformer(\"all-MiniLM-L6-v2\")\n",
    "embeddings = embedding_model.encode(test_sentences, show_progress_bar=True, batch_size=64, device=\"cuda\")\n",
    "\n",
    "print(f\"Generated embeddings of shape: {embeddings.shape}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Create parquet file as historical data source"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%mkdir feature_repo/data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from datetime import datetime, timezone\n",
    "\n",
    "# Create DataFrame\n",
    "df = pd.DataFrame({\n",
    "    \"passage_id\": list(range(len(test_sentences))),\n",
    "    \"passage_text\": test_sentences,\n",
    "    \"embedding\": pd.Series(\n",
    "         [embedding.tolist() for embedding in embeddings],\n",
    "         dtype=object\n",
    "     ),\n",
    "    \"event_timestamp\": [datetime.now(timezone.utc) for _ in test_sentences],\n",
    "})\n",
    "\n",
    "print(\"DataFrame Info:\")\n",
    "print(df.head())\n",
    "print(df[\"embedding\"].apply(lambda x: len(x) if isinstance(x, list) else str(type(x))).value_counts())  # Check lengths\n",
    "\n",
    "# Save to Parquet\n",
    "df.to_parquet(\"feature_repo/data/wiki_dpr.parquet\", index=False)\n",
    "print(\"Saved to wiki_dpr.parquet\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Ensure you are in the feature_repo directory and run feast apply"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%cd feature_repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!feast apply"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Write to the Milvus online store from the parquet file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from feast import FeatureStore\n",
    "import pandas as pd\n",
    "\n",
    "store = FeatureStore(repo_path=\".\")\n",
    "\n",
    "df = pd.read_parquet(\"./data/wiki_dpr.parquet\")\n",
    "store.write_to_online_store(feature_view_name='wiki_passages', df=df)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Define generator model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from transformers import AutoTokenizer, AutoModelForCausalLM, RagConfig, AutoModel\n",
    "\n",
    "generator_model_id = \"ibm-granite/granite-3.2-2b-instruct\"\n",
    "generator_model = AutoModelForCausalLM.from_pretrained(generator_model_id)\n",
    "generator_tokenizer = AutoTokenizer.from_pretrained(generator_model_id)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Initialize Feast Vector Store, Feast Index and FeastRAGRetriever"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import sys\n",
    "sys.path.append(\"..\")\n",
    "from ragproject_repo import wiki_passage_feature_view\n",
    "from feast.vector_store import FeastVectorStore\n",
    "from feast.rag_retriever import FeastIndex, FeastRAGRetriever\n",
    "\n",
    "generator_config=generator_model.config\n",
    "question_encoder = AutoModel.from_pretrained(\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "question_encoder_tokenizer = AutoTokenizer.from_pretrained(\"sentence-transformers/all-MiniLM-L6-v2\")\n",
    "\n",
    "\n",
    "query_encoder_config = {\n",
    "    \"model_type\": \"bert\",\n",
    "    \"hidden_size\": 384\n",
    "}\n",
    "\n",
    "vector_store = FeastVectorStore(\n",
    "    repo_path=\".\",\n",
    "    rag_view=wiki_passage_feature_view,\n",
    "    features=[\"wiki_passages:passage_text\", \"wiki_passages:embedding\", \"wiki_passages:passage_id\"]\n",
    ")\n",
    "\n",
    "feast_index = FeastIndex()\n",
    "\n",
    "config = RagConfig(\n",
    "    question_encoder=query_encoder_config,\n",
    "    generator=generator_config.to_dict(),\n",
    "    index=feast_index\n",
    ")\n",
    "retriever = FeastRAGRetriever(\n",
    "    question_encoder=question_encoder,\n",
    "    question_encoder_tokenizer=question_encoder_tokenizer,\n",
    "    generator_tokenizer=generator_tokenizer,\n",
    "    feast_repo_path=\".\",\n",
    "    feature_view=vector_store.rag_view,\n",
    "    features=vector_store.features,\n",
    "    generator_model=generator_model, \n",
    "    search_type=\"vector\",\n",
    "    id_field=\"passage_id\",\n",
    "    text_field=\"passage_text\",\n",
    "    config=config,\n",
    "    index=feast_index,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "#### Submit a query"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "query = \"What is the capital of Ireland?\"\n",
    "answer = retriever.generate_answer(query, top_k=10)\n",
    "print(\"Generated Answer:\", answer)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": ".venv",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
